package webrtc

import (
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/gofiber/websocket/v2"
	"github.com/pion/rtcp"
	"github.com/pion/webrtc/v3"
)

var (
	RoomsLock sync.RWMutex
	Rooms     map[string]*Room
	Peers     map[*PeerConnectionState]*Room
)

type PeerRole int

const (
	Sender PeerRole = iota
	Receiver
)

var (
	turnConfig = webrtc.Configuration{
		ICETransportPolicy: webrtc.ICETransportPolicyRelay,
		ICEServers: []webrtc.ICEServer{
			{
				URLs: []string{
					"turn:stun.zhikayun.com:3478",
				},
				Username:       "admin",
				Credential:     "123456",
				CredentialType: webrtc.ICECredentialTypePassword,
			},
		},
	}
)

type Room struct {
	RoomId      string
	Rotation    *RotationEvent
	ListLock    sync.RWMutex
	Connections []*PeerConnectionState
	TrackLocals map[string]*webrtc.TrackLocalStaticRTP
}

type PeerConnectionState struct {
	PeerConnection *webrtc.PeerConnection
	Websocket      *ThreadSafeWriter
	Rotation       *RotationEvent
	StopOnce       sync.Once
	Role           PeerRole // 👈 标记是推流者还是接收者

}

type ThreadSafeWriter struct {
	Conn *websocket.Conn
	Send chan []byte
}

func (p *PeerConnectionState) ReadPump() {
	defer func() {
		log.Println("WSClient Close!")
		p.AllStop() // ✅ 统一收尾清理
	}()
	p.Websocket.Conn.SetReadLimit(maxMessageSize)
	p.Websocket.Conn.SetReadDeadline(time.Now().Add(pongWait))
	p.Websocket.Conn.SetPongHandler(func(string) error { p.Websocket.Conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, raw, err := p.Websocket.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error: %v", err)
			}
			log.Printf("error: %v", err)
			break
		}
		fmt.Println("read: " + string(raw))
		var msg RecvMessage
		if err := json.Unmarshal(raw, &msg); err != nil {
			log.Println("Invalid message format:", err)
			continue
		}

		switch msg.Type {
		case ORIENTATION:
			var re RotationEvent
			if err := json.Unmarshal([]byte(msg.Data), &re); err != nil {
				log.Println("Invalid message format:", err)
				continue
			}
			p.Rotation = &re
			room := Peers[p]
			if room == nil {
				continue
			}
			room.Rotation = &re
			for _, connection := range room.Connections {
				if connection.Role == Receiver {
					connection.Websocket.Send <- raw
				}
			}
		case SDP:
			var sd webrtc.SessionDescription
			if err := json.Unmarshal([]byte(msg.Data), &sd); err != nil {
				log.Println("SDP parse error:", err)
				continue
			}
			if err := p.PeerConnection.SetRemoteDescription(sd); err != nil {
				log.Println("SetRemoteDescription error:", err)
			}
		case ICE:
			var candidate webrtc.ICECandidateInit
			if err := json.Unmarshal([]byte(msg.Data), &candidate); err != nil {
				log.Println("ICE parse error:", err)
				continue
			}
			if err := p.PeerConnection.AddICECandidate(candidate); err != nil {
				log.Println("AddICECandidate error:", err)
			}
		case "stop":
			return
		case OPERATE:
			room := Peers[p]
			for _, connection := range room.Connections {
				if connection.Role == Sender {
					connection.Websocket.Send <- raw
				}
			}
		case JOIN:
			feedback := StringMessage{
				Type: FEEDBACK,
				Code: 200,
				Data: "正在连接",
			}
			var realData string
			if err := json.Unmarshal(msg.Data, &realData); err != nil {
				log.Println("解析 data 失败:", err)
				feedback.Data = "投屏码不正确"
				feedback.Code = 500
				val, _ := json.Marshal(feedback)
				p.Websocket.Send <- val
				continue
			}
			if p.Role == Sender {
				err = CreateRoom(realData, p)
				if err != nil {
					feedback.Data = err.Error()
					feedback.Code = 500
					val, _ := json.Marshal(feedback)
					p.Websocket.Send <- val
					continue
				}
				Peers[p].Rotation = p.Rotation
			} else {
				room := Rooms[realData]
				if room == nil {
					log.Println("房间不存在")
					feedback.Data = "房间不存在"
					feedback.Code = 500
					val, _ := json.Marshal(feedback)
					p.Websocket.Send <- val
					continue
				}
				err = WebRTCConn(p, room)
				rotationMsg := room.Rotation
				sendMsg := SendMessage{
					Type: ORIENTATION,
					Data: rotationMsg,
				}
				val, _ := json.Marshal(sendMsg)
				p.Websocket.Send <- val
			}
			if err != nil {
				feedback.Data = err.Error()
				feedback.Code = 500
				val, _ := json.Marshal(feedback)
				p.Websocket.Send <- val
				continue
			}

			// if room, ok := Rooms[inviteCode]; ok {
			// 	err = WebJoinRoom(p, room)
			// 	if err != nil {
			// 		log.Println(err.Error())
			// 		p.StopChan <- struct{}{}
			// 		return
			// 	}
			// 	room.SignalPeerConnections()
			// }
		default:
			log.Println("Unknown message type:", msg.Type)
		}
	}
}

func (p *PeerConnectionState) WritePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		p.AllStop()
	}()
	for {
		select {
		case message, ok := <-p.Websocket.Send:
			p.Websocket.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				return
			}
			fmt.Println("write: " + string(message))
			w, err := p.Websocket.Conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			n := len(p.Websocket.Send)
			for i := 0; i < n; i++ {
				w.Write(newline)
				w.Write(<-p.Websocket.Send)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			p.Websocket.Conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := p.Websocket.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (p *PeerConnectionState) RTCStop() {
	RoomsLock.Lock()
	defer RoomsLock.Unlock()
	if p.PeerConnection != nil {
		p.PeerConnection.Close()
	}
	room := Peers[p]
	if room == nil {
		return
	}
	// 从 room.Connections 中移除当前连接
	newConnections := make([]*PeerConnectionState, 0, len(room.Connections))
	for _, conn := range room.Connections {
		if conn != p {
			newConnections = append(newConnections, conn)
		}
	}
	room.Connections = newConnections
	// 如果房间已经没人了，移除房间
	if len(room.Connections) == 0 {
		delete(Rooms, room.RoomId)
	}
	// 移除当前连接的 Peers 映射
	delete(Peers, p)
	// 再同步其他连接
	// room.SyncConnection()
}

func (p *PeerConnectionState) AllStop() {
	p.StopOnce.Do(func() {
		room := Peers[p]
		p.Websocket.Conn.Close()
		close(p.Websocket.Send)
		p.RTCStop()
		// 如果是callee 通知其他caller重新建立webrtc连接，ws可以不断开
		if p.Role == Sender {
			for idx := range room.Connections {
				if room.Connections[idx].Role == Receiver {
					room.Connections[idx].RTCStop()
				}
			}
		}
	})
}

// Room.AddTrack 修改：
func (p *Room) AddTrack(t *webrtc.TrackRemote) *webrtc.TrackLocalStaticRTP {
	p.ListLock.Lock()
	defer p.ListLock.Unlock()

	trackKey := fmt.Sprintf("%d", t.SSRC()) // 更稳定
	trackLocal, err := webrtc.NewTrackLocalStaticRTP(t.Codec().RTPCodecCapability, t.ID(), t.StreamID())
	if err != nil {
		log.Println("NewTrackLocalStaticRTP error:", err)
		return nil
	}

	p.TrackLocals[trackKey] = trackLocal
	return trackLocal
}

func (p *Room) RemoveTrack(t *webrtc.TrackLocalStaticRTP) {
	p.ListLock.Lock()
	defer p.ListLock.Unlock()

	delete(p.TrackLocals, t.ID())
}

func (p *Room) SyncConnection() {
	p.ListLock.Lock()
	connections := append([]*PeerConnectionState{}, p.Connections...)
	trackLocals := make(map[string]*webrtc.TrackLocalStaticRTP)
	for k, v := range p.TrackLocals {
		trackLocals[k] = v
	}
	p.ListLock.Unlock()
	// fmt.Println(len(connections))
	// for _, conn := range connections {
	p.syncConnection(connections[len(connections)-1], trackLocals)
	// }
}

func (p *Room) syncConnection(conn *PeerConnectionState, trackLocals map[string]*webrtc.TrackLocalStaticRTP) {
	pc := conn.PeerConnection

	// 1. 记录当前 Track ID
	existingSenders := map[string]bool{}
	for _, sender := range pc.GetSenders() {
		if sender.Track() != nil {
			existingSenders[sender.Track().ID()] = true
		}
	}

	// 2. 移除旧的
	for _, sender := range pc.GetSenders() {
		if sender.Track() == nil {
			continue
		}
		if _, ok := trackLocals[sender.Track().ID()]; !ok {
			if err := pc.RemoveTrack(sender); err != nil {
				log.Println("RemoveTrack error:", err)
			}
		}
	}

	// 3. 添加新的
	for trackID, trackLocal := range trackLocals {
		if _, ok := existingSenders[trackID]; !ok {
			if _, err := pc.AddTrack(trackLocal); err != nil {
				log.Println("AddTrack error:", err)
			}
		}
	}

	// 4. Offer/Answer
	offer, err := pc.CreateOffer(nil)
	if err != nil {
		log.Println("CreateOffer error:", err)
		return
	}
	if err := pc.SetLocalDescription(offer); err != nil {
		log.Println("SetLocalDescription error:", err)
		return
	}

	msg, _ := json.Marshal(SendMessage{Type: SDP, Data: offer})

	select {
	case conn.Websocket.Send <- msg:
	case <-time.After(2 * time.Second):
		log.Println("发送超时")
	}

	// go func() {
	// 	time.Sleep(200 * time.Millisecond)
	// 	p.dispatchKeyFrameForConn(conn)
	// }()
}

func (p *Room) dispatchKeyFrameForConn(conn *PeerConnectionState) {
	pc := conn.PeerConnection
	for _, receiver := range pc.GetReceivers() {
		if track := receiver.Track(); track != nil {
			ssrc := track.SSRC()
			if ssrc == 0 {
				log.Println("dispatchKeyFrame: 跳过未初始化的 track（SSRC 为 0）")
				continue
			}
			err := pc.WriteRTCP([]rtcp.Packet{
				&rtcp.PictureLossIndication{MediaSSRC: uint32(ssrc)},
			})
			if err != nil {
				log.Println("发送 PLI 失败:", err)
			}
		}
	}
}

// func (p *Room) SignalPeerConnections() {
// 	p.ListLock.Lock()
// 	defer func() {
// 		p.ListLock.Unlock()
// 		p.DispatchKeyFrame()
// 	}()
// 	attemptSync := func() (tryAgain bool) {
// 		for i := range p.Connections {
// 			// // 👇 加在这里：跳过非接收者（Sender）
// 			// if p.Connections[i].Role != Receiver {
// 			// 	continue
// 			// }
// 			if p.Connections[i].PeerConnection.ConnectionState() == webrtc.PeerConnectionStateClosed {
// 				p.Connections = append(p.Connections[:i], p.Connections[i+1:]...)
// 				log.Println("a", p.Connections)
// 				return true
// 			}

// 			existingSenders := map[string]bool{}
// 			for _, sender := range p.Connections[i].PeerConnection.GetSenders() {
// 				if sender.Track() == nil {
// 					continue
// 				}

// 				existingSenders[sender.Track().ID()] = true

// 				if _, ok := p.TrackLocals[sender.Track().ID()]; !ok {
// 					if err := p.Connections[i].PeerConnection.RemoveTrack(sender); err != nil {
// 						return true
// 					}
// 				}
// 			}

// 			for _, receiver := range p.Connections[i].PeerConnection.GetReceivers() {
// 				if receiver.Track() == nil {
// 					continue
// 				}

// 				existingSenders[receiver.Track().ID()] = true
// 			}

// 			for trackID := range p.TrackLocals {
// 				if _, ok := existingSenders[trackID]; !ok {
// 					if _, err := p.Connections[i].PeerConnection.AddTrack(p.TrackLocals[trackID]); err != nil {
// 						return true
// 					}
// 				}
// 			}

// 			offer, err := p.Connections[i].PeerConnection.CreateOffer(nil)
// 			if err != nil {
// 				return true
// 			}

// 			if err = p.Connections[i].PeerConnection.SetLocalDescription(offer); err != nil {
// 				return true
// 			}
// 			message, err := json.Marshal(SendMessage{
// 				Type: SDP,
// 				Data: offer,
// 			})
// 			if err != nil {
// 				return true
// 			}
// 			p.Connections[i].Websocket.Send <- message
// 		}

// 		return
// 	}

// 	for syncAttempt := 0; ; syncAttempt++ {
// 		if syncAttempt == 25 {
// 			go func() {
// 				time.Sleep(time.Second * 3)
// 				p.SyncAllConnections()
// 			}()
// 			return
// 		}

// 		if !attemptSync() {
// 			break
// 		}
// 	}
// }

func (p *Room) DispatchKeyFrame() {
	p.ListLock.Lock()
	defer p.ListLock.Unlock()

	for i := range p.Connections {
		for _, receiver := range p.Connections[i].PeerConnection.GetReceivers() {
			if receiver.Track() == nil {
				continue
			}

			_ = p.Connections[i].PeerConnection.WriteRTCP([]rtcp.Packet{
				&rtcp.PictureLossIndication{
					MediaSSRC: uint32(receiver.Track().SSRC()),
				},
			})
		}
	}
}

type websocketMessage struct {
	Event string `json:"event"`
	Data  string `json:"data"`
}
